import numpy as np
import matplotlib.pyplot as plt
from target import DuffingOscillator
from flowDMD import FlowDMD, train_flowdmd_xy
from dictionary import DicRBF
from solver import KoopmanGeneralSolver
import argparse
import os
import mindspore as ms
from mindspore import Tensor
from mindspore import save_checkpoint, load_checkpoint, load_param_into_net


# Set context
ms.set_context(mode=ms.PYNATIVE_MODE, device_target='CPU')

duffing_train = DuffingOscillator(n_init=1000, traj_len=10, seed=0)
data_x_train = duffing_train.generate_init_data()
data_y_train = duffing_train.generate_next_data(data_x_train)

duffing_valid = DuffingOscillator(n_init=300, traj_len=10, seed=0)
data_x_valid = duffing_valid.generate_init_data()
data_y_valid = duffing_valid.generate_next_data(data_x_valid)

# Convert to MindSpore tensors with float32
if isinstance(data_x_train, Tensor):
    data_x_train = ops.Cast()(data_x_train, ms.float32)
else:
    data_x_train = Tensor(data_x_train, dtype=ms.float32)

if isinstance(data_y_train, Tensor):
    data_y_train = ops.Cast()(data_y_train, ms.float32)
else:
    data_y_train = Tensor(data_y_train, dtype=ms.float32)

data_train = [data_x_train, data_y_train]
data_valid = [data_x_valid, data_y_valid]


def trainmodel():
    parser = argparse.ArgumentParser(description='FlowDMD训练和预测')
    parser.add_argument('--train', default='True', action='store_true', help='是否进行训练')
    parser.add_argument('--model_path', type=str, default='model.ckpt', help='模型保存路径')
    args = parser.parse_args()

    # 可视化数据
    fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True, figsize=(10, 5))
    
    # Convert tensors to numpy for plotting
    if isinstance(data_x_train, Tensor):
        data_x_plot = data_x_train.asnumpy()
        data_y_plot = data_y_train.asnumpy()
    else:
        data_x_plot = data_x_train
        data_y_plot = data_y_train
    
    ax1.plot(data_x_plot[:, 0], data_x_plot[:, 1], '.')
    ax1.set_xlabel(r'$x_{1}$')
    ax1.set_ylabel(r'$x_{2}$')
    ax2.plot(data_y_plot[:, 0], data_y_plot[:, 1], '.')
    ax2.set_xlabel(r'$x_{1}$')
    ax2.set_ylabel(r'$x_{2}$')

    # 创建模型
    model = FlowDMD(dim=2, num_layers=6, hidden_dim=64, reg=1e-4)
    
    # 检查是否存在模型文件
    if os.path.exists(args.model_path) and not args.train:
        print(f"加载已有模型: {args.model_path}")
        param_dict = load_checkpoint(args.model_path)
        load_param_into_net(model, param_dict)
    else:
        print("开始训练模型...")
        k = train_flowdmd_xy(model, data_xy=data_train, epochs=200, lr=1e-3, alpha=1.0)
        # 保存模型
        print(k)
        save_checkpoint(model, args.model_path)
        print(f"模型已保存到: {args.model_path}")
    
    return model, k if 'k' in locals() else None


def RBFmodel():
    # Convert tensor to numpy for RBF
    if isinstance(data_x_train, Tensor):
        data_x_np = data_x_train.asnumpy()
    else:
        data_x_np = data_x_train
        
    rbf_basis_func = DicRBF(rbf_number=100, regularizer=1e-4)
    rbf_basis_func.build(data_x_np)

    rbf_solver = KoopmanGeneralSolver(dic=rbf_basis_func,
                                     target_dim=2,
                                     reg=0.0)

    # Convert data for RBF solver
    data_train_np = [data_x_np, data_y_train.asnumpy() if isinstance(data_y_train, Tensor) else data_y_train]
    rbf_solver.build(data_train_np)
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5), sharex=False, sharey=True)
    ax2.scatter(rbf_solver.eigenvalues.real, rbf_solver.eigenvalues.imag)
    ax2.set_xlabel(r'Re$(\mu)$')
    ax2.set_title('EDMD (RBF)')

    fig.tight_layout()
    return rbf_solver


def test():
    rbf_solver = RBFmodel()
    model, _ = trainmodel()
    
    traj_len = 50
    duffing_test = DuffingOscillator(n_init=1, traj_len=traj_len, seed=111)
    data_x_test = duffing_test.generate_init_data()
    data_y_test = duffing_test.generate_next_data(data_x_test)
    
    # Convert to float32 tensor
    if isinstance(data_x_test, Tensor):
        data_x_test = ops.Cast()(data_x_test, ms.float32)
    else:
        data_x_test = Tensor(data_x_test, dtype=ms.float32)
        
    if isinstance(data_y_test, Tensor):
        data_y_test = ops.Cast()(data_y_test, ms.float32)
    else:
        data_y_test = Tensor(data_y_test, dtype=ms.float32)
    
    x_traj = data_x_test.asnumpy()

    x0_test = data_x_test[0]
    reshape = ops.Reshape()
    x0_test = reshape(x0_test, (-1, x0_test.shape[-1]))

    # Set model to evaluation mode (equivalent to torch.no_grad())
    model.set_train(False)
    
    # Predictions
    x0_test_np = x0_test.asnumpy()
    x_est_traj_rbf = rbf_solver.predict(x0_test_np, traj_len)
    if isinstance(x_est_traj_rbf, Tensor):
        x_est_traj_rbf = x_est_traj_rbf.asnumpy()
    
    x_est_traj_DL_tensor = model.predict(x0_test, traj_len)
    if isinstance(x_est_traj_DL_tensor, Tensor):
        x_est_traj_DL = x_est_traj_DL_tensor[0].asnumpy()
    else:
        x_est_traj_DL = x_est_traj_DL_tensor[0]
    
    # Calculate errors
    DL_error = np.sqrt(np.mean(np.square(x_est_traj_DL - x_traj)))
    rbf_error = np.sqrt(np.mean(np.square(x_est_traj_rbf - x_traj)))
    
    # Plotting
    fig, (ax1, ax2) = plt.subplots(1, 2, sharey=False, figsize=(12, 5))

    t_grid = np.arange(0, traj_len)
    ax1.plot(t_grid, x_traj[:, 0], 'k-', label='Exact', markevery=2)
    ax1.plot(t_grid, x_est_traj_DL[:, 0], 'bo', label='EDMD-DL', markevery=2)
    ax1.plot(t_grid, x_est_traj_rbf[:, 0], 'gs', label='EDMD (RBF)', markevery=2)

    ax2.plot(t_grid, x_traj[:, 1], 'k-', label='Exact', markevery=2)
    ax2.plot(t_grid, x_est_traj_DL[:, 1], 'bo', label='EDMD-DL', markevery=2)
    ax2.plot(t_grid, x_est_traj_rbf[:, 1], 'gs', label='EDMD (RBF)', markevery=2)

    ax1.set_xlabel(r"$n$")
    ax1.set_ylabel(r"$x_{1}(n)$")
    ax1.legend(loc='best', prop={'size': 14})

    ax2.set_xlabel(r"$n$")
    ax2.set_ylabel(r"$x_{2}(n)$")
    ax2.legend(loc='best', prop={'size': 14})

    print("DL error: ", DL_error,
          "  RBF error: ", rbf_error)
    fig.tight_layout()
    fig.show()


if __name__ == "__main__":
    test() 